import cv2
import numpy as np
from io import BytesIO
import os
import time
from pathlib import Path
import shutil

#此类编写主要是通过一次计算hash的值后，实现多次比对，并通过多线程实现效率的提升；
class compareImg(object):

    #slist为原始的图片的列表， tlist为需要对比的图片列表；
    def __init__(self):
        self.orb = cv2.ORB_create()
        return

    def compORB(self, spath, slist, tlist, drate=0.5):

        tdict = {}
        count = 0
        #先提取tlist中记录的ORG的值
        rdict = {}
        orb = cv2.ORB_create()
        bf = cv2.BFMatcher(cv2.NORM_HAMMING)
        try:        
            for r in tlist:

                kframe = r[0]
                url = r[1]
                groups = r[2]
                kpic = spath+"/"+url              
                img2 = cv2.imread(kpic, cv2.IMREAD_GRAYSCALE)
                # 提取并计算特征点
                kp2, des2 = orb.detectAndCompute(img2, None)
                #只有groups和des2的数据参与到后面的计算中，所以生成的字典可以只包含这两个内容；
                tdict[count] = [des2, groups, kpic]
                count += 1

            for key, val in slist.items():
                # 读取图片
                img1 = cv2.imread(val, cv2.IMREAD_GRAYSCALE)
                kp1, des1 = orb.detectAndCompute(img1, None)
                #对于数据库中保存的关键帧进行比对；
                for k, r in tdict.items():
                    des = r[0]
                    grp = r[1]
                    kp = r[2]
                    # knn筛选结果
                    try:
                        matches = bf.knnMatch(des1, trainDescriptors=des, k=2)   
                        # 查看最大匹配点数目
                        good = [m for (m, n) in matches if m.distance < 0.90 * n.distance]
                        similary = float(len(good))/len(matches)
                        if similary > drate:
                            rdict[key] = grp
                            print("%s, %s, %s, %s" % (val, kp, grp, similary))
                    except:
                        continue
        except Exception as e:
            print('无法计算两张图片相似度', e)
        return rdict
    #比较两张照片的差异点;如果是两张图片相似，则比值趋近于1.0，所以判断是否为小于0.6，则认为其是转场的点;
    def calSimilary(self, pimg1, pimg2):

        try:
            bf = cv2.BFMatcher(cv2.NORM_HAMMING)
            img1 = cv2.imread(pimg1, cv2.IMREAD_GRAYSCALE)
            kp1, des1 = self.orb.detectAndCompute(img1, None)

            img2 = cv2.imread(pimg2, cv2.IMREAD_GRAYSCALE)
            kp2, des2 = self.orb.detectAndCompute(img2, None)
            matches = bf.knnMatch(des1, trainDescriptors=des2, k=2)     
            good = [m for (m, n) in matches if m.distance < 0.95 * n.distance]
            similary = float(len(good))/len(matches)
            return(similary)
        except:
            return 0

    #关键帧比较，前后两帧进行处理，如果差异值大于0.7则认为是转场的图片；
    def keyCompare(self, img_path):

        #准备一个临时目录存放关键帧转场的文件
        tmp_pstr = img_path + "/key"
        tmp_path = Path(tmp_pstr)
        tmp_path.mkdir(parents=True, exist_ok=True)
        #get  path files list;
        dlist = os.listdir(img_path)
        #只获取列表中的.jpg内容；
        nlist = []
        for m in  range(len(dlist)):
            if dlist[m].endswith('.jpg'):
                nlist.append(dlist[m])

        img_list =sorted(nlist)    #文件名按字母排序
        img_nums =len(img_list)
        #先获取第一帧做为种子，来进行处理后续的比对；
        plist = []
        fimg =  img_path + "/" + img_list[0]
        for i in range(img_nums):
            pimg = img_path + "/" +img_list[i]
            sm = self.calSimilary(fimg, pimg)
            if(sm <0.45):
                #把图片拷贝到key目录下
                #先去掉原文件中的后缀，重新组合文件名
                n_f = img_list[i][:-4]
                shutil.copyfile(pimg, tmp_pstr+"/"+n_f+"_"+str(sm)+".jpg")
            fimg = pimg
            
        return self.getPathImgList(tmp_pstr)

    def getPathImgList(self, img_path):
        
        dlist = os.listdir(img_path)
        #只获取列表中的.jpg内容；
        nlist = []
        for m in  range(len(dlist)):
            if dlist[m].endswith('.jpg'):
                nlist.append(dlist[m])

        return nlist
if __name__ == "__main__":
    # p1="/tmp/autostrip/1234567/000346_014802_000053294940.jpg"
    # p2="/tmp/autostrip/1234567/000346_014802_000053294940.jpg"

    comp = compareImg()
    for i in range(1,51):
        comp.keyCompare("/tmp/autostrip/"+str(i))
